#include <Arduino.h>
#include "model_define.hpp"
#include "MnistArray.hpp"


// int16_t example_element[] = {    
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,  38, 254,
//        109,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//         87, 252,  82,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0, 135, 241,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,  45, 244, 150,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,  84, 254,  63,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0, 202, 223,  11,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,  32, 254, 216,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,  95, 254, 195,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 140,
//        254,  77,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//         57, 237, 205,   8,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0, 124, 255, 165,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0, 171, 254,  81,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,  24, 232, 215,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0, 120, 254, 159,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0, 151, 254, 142,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 228, 254,
//         66,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,  61,
//        251, 254,  66,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0, 141, 254, 205,   3,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,  10, 215, 254, 121,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   5, 198, 176,  10,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,
//          0,   0,   0,   0};

void prediction(int16_t example_element[]) {
    
    int input_height = 28;
    int input_width = 28;
    int input_channel = 1;
    int input_exponent = -15;
    int model_input_size = input_height * input_width * input_channel;
    int16_t *model_input = (int16_t *)malloc(model_input_size * sizeof(int16_t));
    int arraySize = sizeof(example_element) / sizeof(example_element[0]);
    // for (int i = 0; i < 784; i++) {
    //   Serial.print(example_element[i]);
    //   Serial.print(" ");
    // }
    // Serial.println();
    // Serial.println("arraySize:"); 
    // Serial.println(arraySize); 


    // 输入数据归一化
    for (int i = 0; i < model_input_size; i++) {
        float normalized_input = example_element[i] / 255.0; // 归一化
        model_input[i] = (int16_t)constrain(normalized_input * (1 << -input_exponent), -32768, 32767);
    }

    Tensor<int16_t> input;
    input.set_element((int16_t *)model_input).set_exponent(input_exponent).set_shape({28, 28, 1}).set_auto_free(false);

    MNIST model;

    dl::tool::Latency latency;

    // model.build(input);

    // model forward
    latency.start();
    // model.build(input)
    model.forward(input);
    latency.end();
    latency.print("MNIST", "forward");

    // 查找最大值的索引
    int16_t *score = model.l7.get_output().get_element_ptr();
    int16_t max_score = score[0];
    int max_index = 0;
    Serial.printf("%d, ", max_score);
    Serial.println("score:");
    for (size_t i = 1; i < 10; i++) {
        Serial.printf("score[%d]:%d, ",i, score[i]);
        if (score[i] > max_score) {
            max_score = score[i];
            max_index = i;
        }
    }
    Serial.printf("\nPrediction Result: %d\n", max_index);
}

void setup() {
  Serial.begin(115200);
  Serial.println("setup!");
  
}

void loop() {
  int16_t dataArray[784];
  int dataCount = 0;
  testSerialInput(dataArray, &dataCount);
  if(dataCount == 784)
  {
    // Serial.printf("\ndataArray size: %d\n", sizeof(dataArray) / sizeof(dataArray[0]));
    // Serial.printf("\ndataCount: %d\n", dataCount);
    prediction(dataArray);
    dataCount = 0;
  }

  // receiveAndPrintArray();


}

